import cv2
import numpy as np

from skimage.exposure import rescale_intensity


def convolve(*, image: np.ndarray, k: np.ndarray):
    ih, iw = image.shape[:2]
    kh, kw = k.shape[:2]
    # 填充
    pad = (kw - 1) // 2
    image = cv2.copyMakeBorder(image, pad, pad, pad, pad,
                               cv2.BORDER_REPLICATE)
    output = np.zeros((ih, iw), dtype='float')

    # 开始扫描
    for y in np.arange(pad, ih + pad):
        for x in np.arange(pad, iw + pad):
            roi = image[y - pad: y + pad + 1, x - pad: x + pad + 1]
            temp = (roi * k).sum()
            output[y - pad, x - pad] = temp
    output: np.ndarray = rescale_intensity(output, in_range=(0, 255))
    output = (output * 255).astype("uint8")
    return output


# 用于图像平滑模糊
small_blur = np.ones([7, 7], dtype='float') / (7 * 7)
big_blur = np.ones([21, 21], dtype='float') / (21 * 21)

# 锐化图像核
sharpen_k = np.array(
    ([0, -1, 0],
     [-1, 5, -1],
     [0, -1, 0]), dtype='int'
)

# 拉普拉斯核进行边缘检测
laplacian_k = np.array(
    ([0, 1, 0],
     [1, -4, 1],
     [0, 1, 0]), dtype='int'
)

# sobel核用于沿X轴和Y轴检测类边缘区域
sobel_k_x = np.array(
    ([-1, 0, 1],
     [-2, 0, 2],
     [-1, 0, 1]), dtype='int'
)
sobel_k_y = np.array(
    ([-1, -2, -1],
     [0, 0, 0],
     [1, 2, 1]), dtype='int'
)

# emboss核
emboss_k = np.array(
    ([-2, -1, 0],
     [-1, 1, 1],
     [0, 1, 2]), dtype='int'
)

# 内核库
k_bank = (
    ('small_blur', small_blur),
    ('big_blur', big_blur),
    ('sharpen_k', sharpen_k),
    ('laplacian_k', laplacian_k),
    ('sobel_k_x', sobel_k_x),
    ('sobel_k_y', sobel_k_y),
    ('emboss_k', emboss_k)
)

if __name__ == '__main__':
    image = cv2.imread('../../dataset/dog_cat/cat.100.jpg')
    # 转换成灰度
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 循环核
    for k_name, k in k_bank:
        print(f'[info]:正在使用核--{k_name}')
        convolve_output = convolve(image=gray, k=k)
        opencv_output = cv2.filter2D(gray, -1, k)

        # 展示核处理后的图片
        cv2.imshow("Original", gray)
        cv2.imshow(f"{k_name} - convolve", convolve_output)
        cv2.imshow(f"{k_name} - opencv", opencv_output)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
